"""
数据处理工具类

提供统一的数据处理功能，包括DataFrame操作、数据清洗、统计计算等。
优化pandas操作，提升数据处理性能。
"""

import pandas as pd
import numpy as np
from typing import Dict, Any, List, Optional, Tuple
import logging
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import seaborn as sns
from io import BytesIO
import base64

logger = logging.getLogger(__name__)


class DataProcessor:
    """
    数据处理器
    
    提供高效的数据处理方法，包括数据清洗、统计计算、可视化等功能。
    """
    
    def __init__(self):
        self.logger = logging.getLogger(self.__class__.__name__)
        # 设置中文字体支持
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
    
    def clean_dataframe(self, df: pd.DataFrame, fill_values: Dict[str, Any] = None) -> pd.DataFrame:
        """
        清洗DataFrame数据
        
        Args:
            df: 需要清洗的DataFrame
            fill_values: 填充值字典
            
        Returns:
            pd.DataFrame: 清洗后的DataFrame
        """
        if df.empty:
            return df
        
        cleaned_df = df.copy()
        
        # 默认填充值
        default_fill_values = {
            'date': pd.NaT,
            'total_steps': 0,
            'total_time': 0,
            'intensity': 0,
            'runtime_hours': 0,
            'battery': 0
        }
        
        if fill_values:
            default_fill_values.update(fill_values)
        
        # 只对存在的列进行填充
        existing_fill_values = {k: v for k, v in default_fill_values.items() if k in cleaned_df.columns}
        cleaned_df = cleaned_df.fillna(existing_fill_values)
        
        return cleaned_df
    
    def format_imei_column(self, df: pd.DataFrame, column: str = 'imei', last_n: int = 6) -> pd.DataFrame:
        """
        格式化IMEI列
        
        Args:
            df: DataFrame
            column: IMEI列名
            last_n: 取最后几位
            
        Returns:
            pd.DataFrame: 格式化后的DataFrame
        """
        if df.empty or column not in df.columns:
            return df
        
        df_copy = df.copy()
        df_copy[column] = df_copy[column].astype(str).str[-last_n:]
        return df_copy
    
    def calculate_attendance_rate(self, df_base: pd.DataFrame, df_work: pd.DataFrame, 
                                total_days: int, imei_col: str = 'imei') -> pd.DataFrame:
        """
        计算出勤率
        
        Args:
            df_base: 基础数据
            df_work: 工作数据
            total_days: 总天数
            imei_col: IMEI列名
            
        Returns:
            pd.DataFrame: 包含出勤率的DataFrame
        """
        self.logger.info("开始计算出勤率")
        
        if df_base.empty:
            return pd.DataFrame()
        
        # 合并数据
        df_merged = pd.merge(df_base, df_work, on=imei_col, how="left")
        df_merged = self.clean_dataframe(df_merged)
        
        # 计算出勤率
        attendance_stats = df_merged.groupby(imei_col).agg(
            工作天数=("date", lambda x: x.nunique())
        ).reset_index()
        
        attendance_stats["总天数"] = total_days
        attendance_stats["出勤率"] = (
            attendance_stats["工作天数"] / attendance_stats["总天数"] * 100
        ).round(2)
        
        # 计算达标情况
        if not df_work.empty and 'total_steps' in df_work.columns and 'total_time' in df_work.columns:
            daily_stats = df_merged.groupby([imei_col, 'date']).agg(
                daily_steps_sum=('total_steps', 'sum'),
                daily_time_sum=('total_time', 'sum')
            ).reset_index()
            
            # 定义达标条件（8000步数和300分钟工作时间）
            daily_stats['达标'] = daily_stats.apply(
                lambda row: 1 if row['daily_steps_sum'] > 8000 and row['daily_time_sum'] > 300 else 0,
                axis=1
            )
            
            compliance_stats = daily_stats.groupby(imei_col).agg(
                达标天数=("达标", lambda x: (x == 1).sum())
            ).reset_index()
            
            # 合并达标情况
            attendance_stats = pd.merge(attendance_stats, compliance_stats, on=imei_col, how='left')
            attendance_stats['达标天数'] = attendance_stats['达标天数'].fillna(0)
            attendance_stats['达标率'] = (
                attendance_stats['达标天数'] / attendance_stats['工作天数'] * 100
            ).round(2)
        
        self.logger.info(f"出勤率计算完成，处理{len(attendance_stats)}条记录")
        return attendance_stats
    
    def calculate_region_statistics(self, df_equipment: pd.DataFrame, df_work: pd.DataFrame, 
                                  total_days: int) -> pd.DataFrame:
        """
        计算区域统计信息
        
        Args:
            df_equipment: 设备数据
            df_work: 工作数据
            total_days: 总天数
            
        Returns:
            pd.DataFrame: 区域统计信息
        """
        self.logger.info("开始计算区域统计信息")
        
        if df_equipment.empty or 'region_name' not in df_equipment.columns:
            return pd.DataFrame()
        
        # 合并数据
        df_merged = pd.merge(df_equipment, df_work, on="imei", how="left")
        df_merged = self.clean_dataframe(df_merged)
        
        # 按区域统计工作天数
        region_work_stats = df_merged.groupby("region_name").agg(
            工作天数=("date", "count"),
        ).reset_index()
        
        # 按区域统计总设备数
        region_total_stats = df_equipment.groupby("region_name").agg(
            设备数量=("imei", "count")
        ).reset_index()
        
        region_total_stats['应出勤天数'] = region_total_stats['设备数量'] * total_days
        
        # 合并统计信息
        region_stats = pd.merge(region_total_stats, region_work_stats, on="region_name", how="left")
        region_stats['工作天数'] = region_stats['工作天数'].fillna(0)
        region_stats['大区出勤率'] = (region_stats['工作天数'] / region_stats['应出勤天数']).round(2)
        
        self.logger.info(f"区域统计完成，处理{len(region_stats)}个区域")
        return region_stats
    
    def identify_no_work_records(self, df_base: pd.DataFrame, df_work: pd.DataFrame,
                               work_indicator_col: str = 'total_steps') -> pd.DataFrame:
        """
        识别无工作记录的员工/设备
        
        Args:
            df_base: 基础数据
            df_work: 工作数据
            work_indicator_col: 工作指示列
            
        Returns:
            pd.DataFrame: 无工作记录的数据
        """
        if df_base.empty:
            return pd.DataFrame()
        
        # 合并数据
        df_merged = pd.merge(df_base, df_work, on="imei", how="left")
        
        # 找出无工作记录的条目
        no_work_mask = df_merged[work_indicator_col].isna()
        no_work_df = df_merged.loc[no_work_mask]
        
        self.logger.info(f"识别出{len(no_work_df)}条无工作记录")
        return no_work_df
    
    def calculate_time_metrics(self, df: pd.DataFrame, time_col: str = 'total_time') -> Dict[str, Any]:
        """
        计算时间相关指标
        
        Args:
            df: 数据框
            time_col: 时间列名
            
        Returns:
            Dict[str, Any]: 时间指标统计
        """
        if df.empty or time_col not in df.columns:
            return {}
        
        # 过滤掉0值
        valid_times = df[df[time_col] > 0][time_col]
        
        if valid_times.empty:
            return {}
        
        # 转换为小时
        hours = valid_times / 60
        
        metrics = {
            'total_hours': hours.sum(),
            'average_hours': hours.mean(),
            'median_hours': hours.median(),
            'max_hours': hours.max(),
            'min_hours': hours.min(),
            'std_hours': hours.std(),
            'records_count': len(valid_times)
        }
        
        return metrics
    
    def format_time_display(self, minutes: float) -> str:
        """
        格式化时间显示（分钟转小时分钟）
        
        Args:
            minutes: 分钟数
            
        Returns:
            str: 格式化的时间字符串
        """
        if pd.isna(minutes) or minutes == 0:
            return "0小时0分钟"
        
        hours = int(minutes // 60)
        remaining_minutes = int(minutes % 60)
        
        return f"{hours}小时{remaining_minutes}分钟"
    
    def calculate_compliance_rate(self, df: pd.DataFrame, threshold_hours: float = 2.0,
                                time_col: str = 'total_time') -> pd.DataFrame:
        """
        计算达标率
        
        Args:
            df: 数据框
            threshold_hours: 达标阈值（小时）
            time_col: 时间列名
            
        Returns:
            pd.DataFrame: 达标率统计
        """
        if df.empty or time_col not in df.columns:
            return pd.DataFrame()
        
        threshold_minutes = threshold_hours * 60
        
        # 按IMEI分组计算达标情况
        compliance_stats = df.groupby('imei').agg(
            总工作天数=('date', 'nunique'),
            达标天数=(time_col, lambda x: (x >= threshold_minutes).sum()),
            总工作时间=(time_col, 'sum')
        ).reset_index()
        
        compliance_stats['达标率'] = (
            compliance_stats['达标天数'] / compliance_stats['总工作天数'] * 100
        ).round(2)
        
        compliance_stats['平均日工作时间'] = compliance_stats['总工作时间'] / compliance_stats['总工作天数']
        
        return compliance_stats
    
    def generate_ranking_table(self, df: pd.DataFrame, metric_col: str, 
                             ascending: bool = False, top_n: int = 5) -> pd.DataFrame:
        """
        生成排名表
        
        Args:
            df: 数据框
            metric_col: 排名依据的列名
            ascending: 是否升序排列
            top_n: 取前n名
            
        Returns:
            pd.DataFrame: 排名表
        """
        if df.empty or metric_col not in df.columns:
            return pd.DataFrame()
        
        # 过滤掉无效值
        valid_df = df[df[metric_col].notna() & (df[metric_col] > 0)]
        
        if valid_df.empty:
            return pd.DataFrame()
        
        # 排序并添加排名
        ranked_df = valid_df.sort_values(metric_col, ascending=ascending).head(top_n).copy()
        ranked_df['排名'] = range(1, len(ranked_df) + 1)
        
        return ranked_df
    
    def create_summary_chart(self, data: Dict[str, float], title: str = "数据摘要") -> str:
        """
        创建摘要图表并返回base64编码
        
        Args:
            data: 数据字典
            title: 图表标题
            
        Returns:
            str: base64编码的图片
        """
        try:
            fig, ax = plt.subplots(figsize=(10, 6))
            
            keys = list(data.keys())
            values = list(data.values())
            
            bars = ax.bar(keys, values, color='skyblue', alpha=0.7)
            ax.set_title(title, fontsize=16, fontweight='bold')
            ax.set_ylabel('数值', fontsize=12)
            
            # 在柱子上添加数值标签
            for bar, value in zip(bars, values):
                height = bar.get_height()
                ax.text(bar.get_x() + bar.get_width()/2., height,
                       f'{value:.1f}', ha='center', va='bottom')
            
            plt.xticks(rotation=45)
            plt.tight_layout()
            
            # 保存为base64
            buffer = BytesIO()
            plt.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
            buffer.seek(0)
            image_base64 = base64.b64encode(buffer.getvalue()).decode()
            plt.close()
            
            return image_base64
            
        except Exception as e:
            self.logger.error(f"创建图表失败: {str(e)}")
            return ""
    
    def optimize_dataframe_memory(self, df: pd.DataFrame) -> pd.DataFrame:
        """
        优化DataFrame内存使用
        
        Args:
            df: 原始DataFrame
            
        Returns:
            pd.DataFrame: 优化后的DataFrame
        """
        if df.empty:
            return df
        
        optimized_df = df.copy()
        
        for col in optimized_df.columns:
            col_type = optimized_df[col].dtype
            
            if col_type != 'object':
                if col_type == 'int64':
                    c_min = optimized_df[col].min()
                    c_max = optimized_df[col].max()
                    if c_min > np.iinfo(np.int8).min and c_max < np.iinfo(np.int8).max:
                        optimized_df[col] = optimized_df[col].astype(np.int8)
                    elif c_min > np.iinfo(np.int16).min and c_max < np.iinfo(np.int16).max:
                        optimized_df[col] = optimized_df[col].astype(np.int16)
                    elif c_min > np.iinfo(np.int32).min and c_max < np.iinfo(np.int32).max:
                        optimized_df[col] = optimized_df[col].astype(np.int32)
                        
                elif col_type == 'float64':
                    c_min = optimized_df[col].min()
                    c_max = optimized_df[col].max()
                    if c_min > np.finfo(np.float32).min and c_max < np.finfo(np.float32).max:
                        optimized_df[col] = optimized_df[col].astype(np.float32)
        
        return optimized_df
    
    def batch_process_dataframes(self, dataframes: Dict[str, pd.DataFrame], 
                               operations: List[str]) -> Dict[str, pd.DataFrame]:
        """
        批量处理多个DataFrame
        
        Args:
            dataframes: DataFrame字典
            operations: 操作列表
            
        Returns:
            Dict[str, pd.DataFrame]: 处理后的DataFrame字典
        """
        processed_dfs = {}
        
        for name, df in dataframes.items():
            processed_df = df.copy()
            
            for operation in operations:
                if operation == 'clean':
                    processed_df = self.clean_dataframe(processed_df)
                elif operation == 'format_imei':
                    processed_df = self.format_imei_column(processed_df)
                elif operation == 'optimize_memory':
                    processed_df = self.optimize_dataframe_memory(processed_df)
            
            processed_dfs[name] = processed_df
        
        return processed_dfs


# 全局数据处理器实例
_global_processor: Optional[DataProcessor] = None


def get_data_processor() -> DataProcessor:
    """获取全局数据处理器实例"""
    global _global_processor
    if _global_processor is None:
        _global_processor = DataProcessor()
    return _global_processor


def set_data_processor(processor: DataProcessor):
    """设置全局数据处理器实例"""
    global _global_processor
    _global_processor = processor